import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import time
import logging

# 加载配置文件
zhipu_api_key = ""
zhipu_api_url = "https://open.bigmodel.cn/api/paas/v4"

# 创建日志记录器
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 定义 Pydantic 模型用于请求参数验证
class VideoRequest(BaseModel):
    prompt: str # 提示语
    with_audio: bool = True # 是否返回音频，默认为 True

# 定义 Pydantic 模型用于响应数据验证
class VideoResponse(BaseModel):
    video_url: str # 视频地址
    cover_image_url: str # 封面图片地址

# 提交文生视频任务的函数
def submit_text_to_video_job(prompt: str, with_audio: bool = True) -> str:
    headers = {
        "Authorization": f"Bearer {zhipu_api_key}",
        "Content-Type": "application/json"
    }
    pyload = {
        "model": "cogvideox-3",
        "prompt": prompt,
        "with_audio": with_audio
    }
    try:
        response = requests.post(zhipu_api_url + "/videos/generations", headers=headers, json=pyload, timeout=300)
        if response.status_code != 200:
            logger.error(f"提交文生视频任务失败，状态码：{response.status_code}，响应内容：{response.text}")
            raise HTTPException(status_code=response.status_code, detail="提交文生视频任务失败")
        submit_response = response.json()
        task_id = submit_response.get("id")
        if not task_id:
            logger.error("提交文生视频任务失败，响应内容：{}".format(submit_response))
            raise HTTPException(status_code=500, detail="提交文生视频任务失败")
        logger.info("提交文生视频任务成功，任务ID：{}".format(task_id))
        return task_id
    except requests.exceptions.RequestException as e:
        logger.error("提交文生视频任务失败，错误信息：{}".format(e))
        raise HTTPException(status_code=500, detail="提交文生视频任务失败")

# 查询文生视频任务状态和结果的函数
def query_text_to_video_job_status(task_id: str):
    headers = {
        "Authorization": f"Bearer {zhipu_api_key}",
        "Content-Type": "application/json"
    }
    status_url = zhipu_api_url + "/async-result/{}".format(task_id)

    start_time = time.time() # 获取当前时间
    logger.info("开始查询文生视频任务状态，任务ID：{}".format(task_id))

    while True:
        try:
            response = requests.get(status_url, headers=headers, timeout=300)
            if response.status_code != 200:
                logger.error("查询文生视频任务状态失败，状态码：{}，响应内容：{}".format(response.status_code, response.text))
                raise HTTPException(status_code=response.status_code, detail="查询文生视频任务状态失败")
            status_response = response.json()
            status = status_response.get("task_status")

            if status == "SUCCESS":
                logger.info("文生视频任务成功，结果：{}".format(status_response))
                video_result = status_response.get("video_result", [])
                if video_result:
                    video_url = video_result[0].get("url")
                    cover_image_url = video_result[0].get("cover_image_url")
                    return {"video_url": video_url, "cover_image_url": cover_image_url}
                else:
                    raise HTTPException(status_code=500, detail="文生视频任务结果为空")

            elif status == "PROCESSING":
                elapsed_time = time.time() - start_time # 计算已用时间
                logger.info("文生视频任务处理中，已用时间：{}秒".format(elapsed_time))
                time.sleep(5) # 等待5秒

            elif status == "FAIL":
                logger.error("文生视频任务状态异常，结果：{}".format(status_response))
                raise HTTPException(status_code=500, detail="文生视频任务状态异常")

            else:
                logger.warning("文生视频任务状态未知，结果：{}".format(status))

        except requests.exceptions.RequestException as e:
            logger.error("查询文生视频任务状态失败，错误信息：{}".format(e))
            time.sleep(5) # 请求失败，等待5秒重试

# 创建 FastAPI 应用
app = FastAPI()

# 提交文生视频任务的接口
@app.post("/zhipuai/video")
async def submit_text_to_video(request: VideoRequest) -> VideoResponse:
    try:
        task_id = submit_text_to_video_job(request.prompt, request.with_audio)
        logger.info("开始处理文生视频任务，任务ID：{}".format(task_id))
        video_result = query_text_to_video_job_status(task_id)
        return VideoResponse(**video_result)
    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error("处理文生视频任务异常，错误信息：{}".format(e))
        raise HTTPException(status_code=500, detail="处理文生视频任务异常")


# 启动 FastAPI 应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)


"""
运行该脚本，启动 FastAPI 应用，并访问 http://localhost:8000/zhipuai/video
POST 请求体中需要包含以下参数：
{
   "prompt": "旺旺小布丁（博美）在草坪上面自由的玩耍",
   "with_audio": false
}
"""
